import json
import time
import re
import os
from openai import OpenAI
from docx import Document

# ========================== 1. 配置区 ==========================
API_KEY = "BearerCubJNNWPCQjooxiEbLGB:IgSlCSEPDRuFXHJlXlMw" 
URL = "https://spark-api-open.xf-yun.com/v2/chat/completions"
INPUT_DOCUMENT_PATH = r"C:/Users/cassi/Desktop/bisai/word-master/豫洛阳-兰郑长干线-CPY-0790-BFGDGS-ZZSYQFGS.docx"
OUTPUT_JSON_PATH = r"C:/Users/cassi/Desktop/bisai/word-master/template_compliance_check.json"
PROCESS_FIRST_CHUNK_ONLY = False  # 仅处理第一块以快速预览效果

# 模板要求内容
TEMPLATE_REQUIREMENTS = {
    # 按照用户提供的“满分说明”来定义必备要点
    '管道本体管控措施': [
        '外检测', '内检测', '环焊缝排查', '电位测试', '阴极保护',
        '检测单位', '检测时间', '结果分析'
    ],
    '外部环境风险管控': [
        '人防', '物防', '技防', '巡护频次', '三桩一牌',
        '视频监控', 'SCADA', '管廊带通透'
    ],
    '事故状态下前期处置': [
        '社会依托资源', '应急通讯录', '联动机制'
    ]
}

# ========================== 2. 功能函数区 ==========================
def get_answer(message):
    try:
        client = OpenAI(
            api_key="sk-ma0HIkN3beK0mtThtDzIfRLN5IzpL4agmB9RSbpHk53h7BF5",
            base_url="https://www.chataiapi.com/v1",
        )
        response = client.chat.completions.create(
            model="gemini-2.5-pro",
            messages=message,
            temperature=0.7,
        )
        content = response.choices[0].message.content or ""
        if content.strip():
            print("\n\n*******************模型生成内容如下********************\n")
            print(content)
        return content
    except Exception as e:
        print(f"API请求失败：{str(e)}")
        return ""

def read_document(file_path):
    """
    读取 Word 文件，提取所有文本内容。
    """
    try:
        doc = Document(file_path)
        full_text = ""
        
        # 提取段落文本
        for paragraph in doc.paragraphs:
            if paragraph.text.strip():
                full_text += paragraph.text.strip() + "\n"
        
        # 提取表格文本
        for table in doc.tables:
            for row in table.rows:
                for cell in row.cells:
                    if cell.text.strip():
                        full_text += cell.text.strip() + " | "
                full_text += "\n"
        
        return full_text.strip()
    except Exception as e:
        print(f"Word文档读取失败：{str(e)}")
        return ""

def semantic_chunker(text, max_chars=250, overlap=50):
    """将长文本进行自然的语义切块"""
    text = re.sub(r'\n\s*\n', '\n', text)
    paragraphs = text.split('\n')
    chunks = []
    current_chunk = ""
    
    for p in paragraphs:
        p = p.strip()
        if not p:
            continue
        if len(current_chunk) + len(p) + 1 <= max_chars:
            current_chunk += p + "\n"
        else:
            if current_chunk:
                chunks.append(current_chunk.strip())
            current_chunk = p + "\n"
    if current_chunk:
        chunks.append(current_chunk.strip())
    return chunks

def extract_section_content(text, section_name):
    """提取特定章节的内容"""
    section_content = []
    
    # 构建章节匹配模式
    section_patterns = [
        rf'{section_name}[^\n]*\n(.*?)(?=\n\s*\d+\.\d+|\n\s*[一二三四五六七八九十]+[、.]|\n\s*第[一二三四五六七八九十]+|\Z)',
        rf'{section_name}[：:](.*?)(?=\n\s*\d+\.\d+|\n\s*[一二三四五六七八九十]+[、.]|\Z)',
        rf'.*{section_name}.*?\n(.*?)(?=\n.*(?:管控|措施|处置|方案).*\n|\Z)'
    ]
    
    for pattern in section_patterns:
        matches = re.finditer(pattern, text, re.DOTALL | re.IGNORECASE)
        for match in matches:
            content = match.group(1).strip()
            if len(content) > 50:  # 过滤过短的内容
                section_content.append({
                    'content': content,
                    'start_pos': match.start(),
                    'end_pos': match.end(),
                    'match_pattern': pattern
                })
    
    return section_content


# ========================== 核心：模板合规性负样本提示词模板 ==========================
def get_negative_prompt(text_chunk, source_label, prompt_type):
    prompt_templates = {
        1: f"""任务：为【{source_label}】生成“管道本体管控措施不完整”负样本。
要求：
1. 仅修改管道本体管控措施相关内容，删除或弱化以下要点：外检测、内检测、环焊缝排查、电位测试、阴极保护、检测单位、检测时间、结果分析
2. 保持其他章节内容不变，仅让管道本体管控措施显得不完整或模糊
3. 文本需保持语义通顺与上下文连贯，但显著缺少上述关键要点
原始文本：
---
{text_chunk}
---
输出：仅返回修改后的文本，无需额外解释。""",
        
        2: f"""任务：为【{source_label}】生成“外部环境风险管控不符合模板”负样本。
要求：
1. 仅修改外部环境风险管控相关内容，删除或弱化以下要点：人防、物防、技防、巡护频次、三桩一牌、视频监控、SCADA、管廊带通透
2. 保持其他章节内容不变，使环境风险管控表述不具体、不成体系或不符合模板
3. 文本需保持语义通顺与上下文连贯，但明显缺失关键管控维度
原始文本：
---
{text_chunk}
---
输出：仅返回修改后的文本，无需额外解释。""",
        
        3: f"""任务：为【{source_label}】生成“事故状态下前期处置描述不规范”负样本。
要求：
1. 仅修改事故状态下前期处置相关内容，删除或弱化以下要点：社会依托资源、应急通讯录、联动机制
2. 保持其他章节内容不变，使应急处置描述缺乏必要联动与资源支撑
3. 文本需保持语义通顺与上下文连贯，但明显不符合处置规范
原始文本：
---
{text_chunk}
---
输出：仅返回修改后的文本，无需额外解释。"""
    }
    return prompt_templates.get(prompt_type, "")

# ========================== 主程序 ==========================
if __name__ == '__main__':
    document_paths = [INPUT_DOCUMENT_PATH]
    
    try:
        file_name_for_prompt = os.path.basename(INPUT_DOCUMENT_PATH)
        print(f"准备处理文档: {file_name_for_prompt}")
    except:
        file_name_for_prompt = "高后果区管控方案"
        print(f"准备处理文档: {INPUT_DOCUMENT_PATH}")
    
    all_generated_data = []  # 存储「相同正样本+单个负样本」的独立字典
    for doc_path in document_paths:
        print(f"\n{'='*80}\n正在处理文件: {doc_path}\n{'='*80}")
        
        # 提取Word文档文本
        full_text = read_document(doc_path)
        if not full_text:
            print("Word文档读取失败")
            continue
        
        print(f"文件读取成功，总字数: {len(full_text)}")
        text_chunks = semantic_chunker(full_text)
        print(f"文本被切分为 {len(text_chunks)} 个语义块进行处理。\n")
        if PROCESS_FIRST_CHUNK_ONLY and text_chunks:
            text_chunks = text_chunks[:1]
            print("已开启快速预览：仅处理第 1 个文本块。\n")
        
        # 遍历每个文本块，生成3个负样本（循环1-3次）
        for chunk_idx, chunk in enumerate(text_chunks):
            print(f"--- 正在处理第 {chunk_idx+1}/{len(text_chunks)} 个文本块 ---")
            source_label = f"原始文本 ({file_name_for_prompt} 文本块{chunk_idx+1})"
            
            # 关键：循环1-3次，对应3类模板合规性错误
            for prompt_type in range(1, 4):
                neg_type = [
                    "管道本体管控措施不完整",
                    "外部环境风险管控不符合模板", 
                    "事故状态下前期处置描述不规范"
                ][prompt_type-1]
                print(f"  → 生成第{prompt_type}/3个负样本（类型：{neg_type}）")
                
                # 获取提示词并调用模型
                prompt = get_negative_prompt(chunk, source_label, prompt_type)
                chat_message = [{"role": "user", "content": prompt}]
                negative_sample = get_answer(chat_message)
                
                # 生成目标格式：每个负样本1个独立字典
                if negative_sample.strip():
                    single_dict = {
                        "messages": [
                            {
                                "content": negative_sample.strip(),
                                "role": "user"
                            },
                            {
                                "content": f"检测到模板合规性问题：{neg_type}。正确的描述应该是：\n{chunk}",
                                "role": "assistant"
                            }
                        ],
                        "images": []
                    }
                    all_generated_data.append(single_dict)
                else:
                    print(f"  → 第{prompt_type}个负样本生成失败，跳过")
                
                time.sleep(1)  # 避免API限流
    
    # 写入输出文件
    if all_generated_data:
        output_dir = os.path.dirname(OUTPUT_JSON_PATH)
        if not os.path.exists(output_dir):
            try:
                os.makedirs(output_dir)
                print(f"已创建输出目录: {output_dir}")
            except OSError as e:
                print(f"错误: 创建输出目录失败: {e}")
                all_generated_data = []
        
        if all_generated_data:
            try:
                with open(OUTPUT_JSON_PATH, "w", encoding="utf-8") as f:
                    json.dump(all_generated_data, f, ensure_ascii=False, indent=4)
                print(f"\n\n{'='*80}")
                print(f"✅ 处理完成！数据保存至：{OUTPUT_JSON_PATH}")
                print(f"共生成 {len(all_generated_data)} 个对话样本（每个包含负样本user + 正样本assistant）")
                print(f"配置：3种类型 × {len(text_chunks)}个文本块 = {3 * len(text_chunks)} 个对话")
                print("="*80 + "\n")
            except Exception as e:
                print(f"错误: 写入JSON失败: {e}")
    else:
        print("\n\n处理结束，无生成数据。请检查PDF文件或API配置。")
